# coding: utf-8
from __future__ import print_function
import argparse

import torch
from torch.autograd import Variable

import logging
import pickle
import codecs
import numpy as np
import nltk
import traceback
import os


def parse_args():
    parser = argparse.ArgumentParser()
    # Model parameters.
    parser.add_argument('--model_dir', type=str, default='./models/primary/')  # dictionary的地址
    parser.add_argument('--cuda', action='store_true',
                        help='use CUDA')
    args = parser.parse_args()
    args.dict_path = os.path.join(args.model_dir, 'dict.pkl')
    args.model_path = os.path.join(args.model_dir, 'model.pt')
    return args



class Client(object):

    def __init__(self, args):
        with open(args.model_path, 'rb') as f:
            self.model = torch.load(f)
        self.model.eval()

        if args.cuda:
            self.model.cuda()
        else:
            self.model.cpu()

        logging.info('use dictionary: {}'.format(args.dict_path))
        with open(args.dict_path, 'rb') as f:
            self.dictionary = pickle.load(f)
        self.ntokens = len(self.dictionary)
        self.zero_hidden = self.model.init_hidden(1)
        self.clear()

    def clear(self):
        self.word_logit = None
        self.word_logp = None
        self.hidden = self.zero_hidden
        self.history = []
        self.forward('_bos')

    def make_input(self, idx):
        return Variable(
            torch.LongTensor([[idx]]),
            volatile=True
        )

    def parse_sentence(self, sentence):
        # 返回feed的word_token, word_id
        sentence = sentence.lower()
        sentence = sentence.replace('<unk>', '_unk')  # 统一改成这种表达
        word_tokens = nltk.tokenize.word_tokenize(sentence)
        feed_word = []
        feed_idx = []
        for word in word_tokens:
            if word in self.dictionary.word2idx:
                feed_idx.append(self.dictionary.word2idx[word])
                feed_word.append(word)
            else:
                feed_idx.append(self.dictionary.word2idx['_unk'])
                feed_word.append('_unk')
        return feed_word, feed_idx

    def forward(self, sentence):
        # forward 一个sentence，改变self.hidden
        feed_word, feed_idx = self.parse_sentence(sentence)
        if '_bos' in feed_word:
            assert feed_word == ['_bos']  # 只允许_bos被单独传递进来，此时没有feed_time_logp输出
            no_feed_time_logp = True
        else:
            no_feed_time_logp = False
        feed_time_logp = []
        for i, idx in enumerate(feed_idx):
            # 真正将hidden给改变（外界对hidden无兴趣）
            if not no_feed_time_logp:
                feed_time_logp.append(self.word_logp[idx])
                self.history.append(feed_word[i])
            output, self.hidden = self.model(self.make_input(idx), self.hidden)
            self.word_logit = output.squeeze().data.cpu().numpy()  # output就是logit
            self.word_logp = self.word_logit - self.word_logit.max()  # 近似
        return feed_word, feed_idx, feed_time_logp

    def sim_forward(self, sentence):
        # 模拟foward, 不改变self.hidden
        feed_word, feed_idx = self.parse_sentence(sentence)
        feed_time_logp = []
        now_hidden = self.hidden  # 只用这一次
        word_logit = self.word_logit
        word_logp = self.word_logp
        for idx in feed_idx:
            feed_time_logp.append(word_logp[idx])
            output, now_hidden = self.model(self.make_input(idx), now_hidden)
            word_logit = output.squeeze().data.cpu().numpy()  # output就是logit
            word_logp = word_logit - word_logit.max()  # 近似
        return feed_word, feed_idx, feed_time_logp

    def clear_and_score(self, sentence):
        # 先clear，然后计算一个句子的平均logp
        # 结束后hidden为_bos的状态
        self.clear()
        feed_word, feed_idx = self.parse_sentence(sentence)
        feed_time_logp = []
        now_hidden = self.hidden  # 只用这一次
        word_logit = self.word_logit
        word_logp = self.word_logp
        for idx in feed_idx:
            feed_time_logp.append(word_logp[idx])
            output, now_hidden = self.model(self.make_input(idx), now_hidden)
            word_logit = output.squeeze().data.cpu().numpy()  # output就是logit
            word_logp = word_logit - word_logit.max()  # 近似
        mean_logp = np.array(feed_time_logp).mean()
        return feed_word, feed_idx, feed_time_logp, mean_logp

    def state(self, top_n=5):
        index = np.argsort(self.word_logp)[::-1]
        index = index[:top_n]
        logp = self.word_logp[index]
        words = map(self.dictionary.to_word, list(index))
        result = zip(words, list(logp))
        return result


def _f(float_num):
    float_num = float(float_num)
    return '{:3f}'.format(float_num)


class InteractiveClient(Client):

    def __init__(self, args):
        super(InteractiveClient, self).__init__(args)

    def run(self):
        while True:
            try:
                input = raw_input('>> ')
                self.parse_input(input)
            except Exception:
                traceback.print_exc()

    def print_state(self):
        print('--------history-----')
        print(' '.join(self.history))
        print('---------next-------')
        state = self.state()
        for line in state:
            print('word: {:^10} | logp:{:^10}'.format(line[0], _f(line[1])))

    def print_feed(self, feed_word, feed_idx, feed_time_logp):
        for i in range(len(feed_word)):
            print('feeded word: {:^10} | feed time logp: {:^10}'.format(feed_word[i], _f(feed_time_logp[i])))

    def parse_input(self, input):
        input = input.strip()
        if input.find(' ') != -1:
            command = input[:input.find(' ')]
            content = input[input.find(' ') + 1:].strip()
        else:
            command = input.strip()
            content = ' '
        if command == 'p':
            self.print_state()
        elif command == 'f':
            self.print_feed(*self.forward(content))
            self.print_state()
        elif command == 'cf':
            self.clear()
            print('---------reset-------')
            self.print_feed(*self.forward(content))
            self.print_state()
        elif command == 't':
            sents = content.split(';')
            for sent in sents:
                sent = sent.strip()
                if sent == '':
                    continue
                print('-------------------')
                self.print_feed(*self.forward(sent))
            print('-------------------')
        elif command == 'c':
            self.clear()
            print('reseted!')
        elif command == 's':
            self.clear()
            feed_word, feed_idx, feed_time_logp, mean_logp = self.clear_and_score(content)
            self.print_feed(feed_word, feed_idx, feed_time_logp)
            print('feeded: {}  ||   score: {:^10}'.format(' '.join(feed_word), _f(mean_logp)))
        else:
            raise NotImplementedError


if __name__ == '__main__':
    args = parse_args()
    c = InteractiveClient(args)
    c.run()
